import { BaseMessage } from "../message/base";
import { get_encoding, encoding_for_model, type Tiktoken, type TiktokenModel } from "tiktoken";

export class TokenStatistical {
    private encodeCache: Record<string, Tiktoken> = {};
    private static instance: TokenStatistical;

    private constructor() {}

    public static getInstance(): TokenStatistical {
        if (!TokenStatistical.instance) {
            TokenStatistical.instance = new TokenStatistical();
        }
        return TokenStatistical.instance;
    }

    private getEncode(modelType?: string): Tiktoken {
        if (modelType) {
            if (!this.encodeCache[modelType]) {
                try {
                    this.encodeCache[modelType] = encoding_for_model(modelType as TiktokenModel);
                    return this.encodeCache[modelType];
                } catch (_error) {
                    // TODO: 日志记录
                    // console.log(`tiktoken encoding for Model type ${modelType} is not supported`);
                }
            }
        }
        if (!this.encodeCache["cl100k_base"]) {
            this.encodeCache["cl100k_base"] = get_encoding("cl100k_base");
        }
        return this.encodeCache["cl100k_base"];
    }

    /**
     * 获取token数量
     * @param message 消息
     * @param modelType 模型类型，留空则使用cl100k_base进行统计
     * @returns token数量
     */
    public getTokenNumberFromMessage(messages: BaseMessage[] | string, modelType?: string): number {
        const encode = this.getEncode(modelType);
        if (typeof messages === "string") {
            return encode.encode(messages).length;
        }
        const message = messages.reduce((acc, cur) => acc + cur.content, "");
        return encode.encode(message).length;
    }
}
